# 数据库连接与会话管理模块
"""
该模块负责数据库连接的创建、自定义会话类的定义以及FastAPI依赖项的实现。
主要功能包括：
1. 创建数据库引擎
2. 实现带权限控制的数据库会话类
3. 提供FastAPI依赖项用于获取数据库会话
"""

from typing import Annotated

from fastapi import Depends, HTTPException, Request
from sqlmodel import Session, create_engine
from sqlmodel import select as _select
from starlette import status

from app.core.config import settings

# 创建数据库引擎，使用配置文件中的数据库URL和引擎选项
engine = create_engine(
    settings.database_config.url,
    pool_pre_ping=True,
    pool_recycle=300,
)


# Code below omitted 👇


class SessionBackend(Session):
    """
    自定义数据库会话类，扩展了SQLAlchemy的Session类，增加了用户权限控制功能。
    """

    def __init__(self, *args, **kwargs):
        """
        初始化会话实例
        """
        super().__init__(*args, **kwargs)
        self.user = None  # 当前会话关联的用户对象

    def set_user(self, user):
        """
        设置当前会话的用户对象
        """
        self.user = user

    def has_permission(self, code: str) -> bool:
        """
        检查当前用户是否具有指定权限码的权限
        """
        # 检查用户是否已登录
        if not self.user:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED, detail="用户未登录"
            )

        # 检查用户是否具有指定权限
        for permission in self.user.permissions:
            if permission.code == code:
                return True

        return False

    @staticmethod
    def _extract_table_names_from_entity(entity):
        """
        从查询实体中提取表名集合

        该方法用于解析SQL查询中的实体对象，提取其对应的数据库表名。支持多种类型的实体：
        1. SQLModel模型类（具有__tablename__属性）
        2. 表对象（具有table和name属性）
        3. 复合查询对象（具有clauses属性）

        Args:
            entity: 查询实体对象，可以是模型类、表对象或复合查询对象

        Returns:
            Set[str]: 包含所有相关表名的集合，用于后续权限检查
        """
        table_names = set()
        if hasattr(entity, "__tablename__"):
            table_names.add(getattr(entity, "__tablename__"))
        elif hasattr(entity, "table") and hasattr(entity.table, "name"):
            table_names.add(getattr(entity.table, "name"))
        elif hasattr(entity, "clauses"):
            for clause in entity.clauses.clauses:
                if hasattr(clause, "table") and hasattr(clause.table, "name"):
                    table_names.add(getattr(clause.table, "name"))
        return table_names

    def _check_read_permission_for_tables(self, table_names):
        """
        检查用户对指定表是否具有读取权限

        遍历所有涉及的表名，逐一检查当前用户是否具有对应表的读取权限。
        权限检查基于"表名 can read"的权限码格式进行。

        Args:
            table_names (Set[str]): 需要检查权限的表名集合

        Raises:
            HTTPException: 当用户对任一表没有读取权限时抛出401未授权异常
        """
        if self.user.is_admin:
            return
        for table_name in table_names:
            if self.has_permission(f"table_{table_name}_!read"):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=f"{table_name} 表无法读取数据, 当前用户权限不足",
                )

    def select(self, *entities, **__kw):
        """
        重写select方法，在查询数据前进行权限检查, 所有关联表都必须要有权限

        该方法在执行SQL查询前，会先分析查询涉及的所有表，然后检查当前用户是否对这些表
        具有读取权限。只有当用户对所有涉及的表都有读取权限时，才允许执行查询。

        Args:
            *entities: 查询涉及的实体对象，可以是模型类、表对象等
            **__kw: 其他关键字参数，支持skip_permission_check参数跳过权限检查

        Returns:
            Select: SQLAlchemy的Select对象，可用于执行查询

        Raises:
            HTTPException: 当用户没有足够权限时抛出401未授权异常
        """
        # 是否跳过检查, 必须使用pop
        if __kw.pop("skip_permission_check", False):
            # 调用_select方法构造查询对象
            return _select(*entities)

        all_table_names = set()
        for entity in entities:
            table_names = self._extract_table_names_from_entity(entity)
            all_table_names.update(table_names)

        self._check_read_permission_for_tables(all_table_names)

        # 调用_select方法构造查询对象
        stmt = _select(*entities)
        return stmt

    def delete(self, instance):
        """
        重写删除方法，在删除数据前进行权限检查
        """
        if self.user.is_admin:
            return super().delete(instance)
        # 获取实体对象的表名
        if not (table_name := getattr(instance, "__tablename__")):
            raise ValueError(
                "删除的对象必须是 SQLModel 模型实例, 且必须指定__tablename__"
            )

        # 检查用户是否具有删除该表数据的权限
        if self.has_permission(f"table_{table_name}_delete"):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"{table_name} 表无法删除数据, 当前用户权限不足",
            )

        # 调用父类的delete方法
        return super().delete(instance)

    def add(self, instance, _warn: bool = True):
        """
        重写添加方法，在添加或更新数据前进行权限检查
        """
        if self.user.is_admin:
            return super().add(instance, _warn)
        # 获取实体对象的表名
        if not (table_name := getattr(instance, "__tablename__")):
            raise ValueError(
                "添加的对象必须是 SQLModel 模型实例, 且必须指定__tablename__"
            )

        # 检查对象是否是更新操作（通过是否有主键值判断）
        is_update = False
        if hasattr(instance, "id") and getattr(instance, "id", None) is not None:
            is_update = True

        if is_update:
            # 更新操作，检查用户是否具有更新该表数据的权限
            if self.has_permission(f"table_{table_name}_!update"):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=f"{table_name} 表无法更新数据, 当前用户权限不足",
                )
        else:
            # 添加操作，检查用户是否具有添加该表数据的权限
            if self.has_permission(f"table_{table_name}_!add"):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail=f"{table_name} 表无法添加数据, 当前用户权限不足",
                )

        # 调用父类的add方法
        return super().add(instance, _warn=_warn)


def get_db(request: Request):
    """
    FastAPI依赖项，用于获取数据库会话。
    确保每个请求都有独立的数据库会话，并在请求结束时正确关闭。
    """
    # 如果请求状态中还没有数据库会话，则创建一个新的
    if not hasattr(request.state, "db"):
        request.state.db = SessionBackend(engine)
    # 获取数据库会话
    db = request.state.db
    try:
        # 提供数据库会话给依赖注入系统
        yield db
    finally:
        # 确保在请求结束时关闭数据库会话
        db.close()


SessionDep = Annotated[SessionBackend, Depends(get_db)]
